#!/bin/bash

# 配置文件路径
CONFIG_FILE="$HOME/.ssh/tunnel_config"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# 加载配置
load_config() {
    if [ -f "$CONFIG_FILE" ]; then
        source "$CONFIG_FILE"
    else
        setup_config
    fi
}

# 配置向导
setup_config() {
    echo -e "${YELLOW}首次运行，请配置连接信息${NC}"
    read -p "请输入服务器用户名: " SERVER_USER
    read -p "请输入服务器IP地址: " SERVER_IP
    read -p "请输入转发端口(默认10022): " SERVER_PORT
    SERVER_PORT=${SERVER_PORT:-10022}
    
    # 保存配置
    mkdir -p "$(dirname "$CONFIG_FILE")"
    cat > "$CONFIG_FILE" << EOF
SERVER_USER="$SERVER_USER"
SERVER_IP="$SERVER_IP"
SERVER_PORT="$SERVER_PORT"
EOF
    
    echo -e "${GREEN}配置已保存${NC}"
}

# 检查依赖
check_dependencies() {
    local deps=("ssh" "nc" "autossh")
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            echo -e "${RED}缺少依赖: $dep${NC}"
            echo "正在安装..."
            sudo apt-get update && sudo apt-get install -y "$dep"
        fi
    done
}

# 检查SSH密钥
setup_ssh_key() {
    if [ ! -f "$HOME/.ssh/id_rsa" ]; then
        echo -e "${YELLOW}未找到SSH密钥，正在生成...${NC}"
        ssh-keygen -t rsa -N "" -f "$HOME/.ssh/id_rsa"
        echo -e "${GREEN}密钥生成完成${NC}"
    fi
    
    echo -e "${YELLOW}是否需要将SSH密钥复制到服务器？(y/n)${NC}"
    read -r answer
    if [ "$answer" = "y" ]; then
        ssh-copy-id "$SERVER_USER@$SERVER_IP"
    fi
}

# 检查连接状态
check_connection() {
    nc -z "$SERVER_IP" 22 >/dev/null 2>&1
    return $?
}

# 创建隧道
# 添加服务器配置检查函数
check_server_config() {
    echo -e "${YELLOW}正在检查服务器配置...${NC}"
    ssh "$SERVER_USER@$SERVER_IP" 'grep -E "GatewayPorts|AllowTcpForwarding" /etc/ssh/sshd_config' > /dev/null 2>&1
    if [ $? -ne 0 ]; then
        echo -e "${RED}服务器SSH配置可能不正确，请确保有以下配置：${NC}"
        echo -e "GatewayPorts yes"
        echo -e "AllowTcpForwarding yes"
        echo -e "\n${YELLOW}是否要自动配置服务器？(y/n)${NC}"
        read -r answer
        if [ "$answer" = "y" ]; then
            ssh "$SERVER_USER@$SERVER_IP" "sudo bash -c 'echo \"GatewayPorts yes\" >> /etc/ssh/sshd_config && echo \"AllowTcpForwarding yes\" >> /etc/ssh/sshd_config && systemctl restart sshd'"
            echo -e "${GREEN}服务器配置已更新${NC}"
        fi
    fi
}

# 获取当前用户名
LOCAL_USER=$(whoami)

# 修改显示连接信息函数
show_connection_info() {
    echo -e "\n${GREEN}=== 连接信息 ===${NC}"
    echo -e "服务器地址: ${YELLOW}${SERVER_IP}${NC}"
    echo -e "服务器用户: ${YELLOW}${SERVER_USER}${NC}"
    echo -e "转发端口: ${YELLOW}${SERVER_PORT}${NC}"
    echo -e "本地SSH端口: ${YELLOW}22${NC}"
    
    echo -e "\n${GREEN}=== 连接状态检查 ===${NC}"
    # 检查服务器SSH是否可访问
    if nc -z -w5 ${SERVER_IP} 22 2>/dev/null; then
        echo -e "服务器SSH状态: ${GREEN}可访问${NC}"
    else
        echo -e "服务器SSH状态: ${RED}不可访问${NC}"
    fi
    
    # 检查转发端口是否开放
    if nc -z -w5 ${SERVER_IP} ${SERVER_PORT} 2>/dev/null; then
        echo -e "转发端口状态: ${GREEN}已开放${NC}"
    else
        echo -e "转发端口状态: ${RED}未开放${NC}"
    fi
    
    # 检查隧道进程
    if pgrep -f "autossh.*${SERVER_PORT}" >/dev/null; then
        echo -e "隧道进程状态: ${GREEN}运行中${NC}"
    else
        echo -e "隧道进程状态: ${RED}未运行${NC}"
    fi
    
    echo -e "\n${GREEN}=== 连接步骤 ===${NC}"
    echo -e "1. 确保隧道显示"运行中""
    echo -e "2. 确保转发端口显示"已开放""
    echo -e "3. 使用以下命令连接："
    echo -e "   ${YELLOW}ssh -p ${SERVER_PORT} ${LOCAL_USER}@${SERVER_IP}${NC}"
    
    echo -e "\n${GREEN}=== 故障排除 ===${NC}"
    echo -e "1. 如果隧道未运行：选择菜单选项 1 启动隧道"
    echo -e "2. 如果端口未开放：检查服务器防火墙设置"
    echo -e "3. 如果配置有误：确保服务器 /etc/ssh/sshd_config 包含："
    echo -e "   ${YELLOW}GatewayPorts yes"
    echo -e "   AllowTcpForwarding yes${NC}"
    echo -e "4. 配置修改后需要重启SSH服务："
    echo -e "   ${YELLOW}sudo systemctl restart sshd${NC}"
    echo -e "5. 如果仍然无法连接：查看日志（选项 5）"
}

# 修改创建隧道函数
create_tunnel() {
    echo -e "${YELLOW}正在检查服务器配置...${NC}"
    check_server_config
    
    echo -e "${YELLOW}正在创建隧道...${NC}"
    autossh -M 0 -fNR "$SERVER_PORT":localhost:22 "$SERVER_USER@$SERVER_IP" \
        -o ServerAliveInterval=60 \
        -o ServerAliveCountMax=3 \
        -o ExitOnForwardFailure=yes

    # 等待隧道建立
    sleep 3
    if nc -z -w5 ${SERVER_IP} ${SERVER_PORT} 2>/dev/null; then
        echo -e "${GREEN}隧道创建成功${NC}"
    else
        echo -e "${RED}隧道创建失败，请检查服务器配置和防火墙设置${NC}"
    fi
}

# 显示菜单
# 显示连接信息
show_connection_info() {
    echo -e "\n${GREEN}=== 连接信息 ===${NC}"
    echo -e "服务器地址: ${YELLOW}${SERVER_IP}${NC}"
    echo -e "服务器用户: ${YELLOW}${SERVER_USER}${NC}"
    echo -e "转发端口: ${YELLOW}${SERVER_PORT}${NC}"
    echo -e "本地SSH端口: ${YELLOW}22${NC}"
    
    echo -e "\n${GREEN}=== 连接状态检查 ===${NC}"
    # 检查服务器SSH是否可访问
    if nc -z -w5 ${SERVER_IP} 22 2>/dev/null; then
        echo -e "服务器SSH状态: ${GREEN}可访问${NC}"
    else
        echo -e "服务器SSH状态: ${RED}不可访问${NC}"
    fi
    
    # 检查转发端口是否开放
    if nc -z -w5 ${SERVER_IP} ${SERVER_PORT} 2>/dev/null; then
        echo -e "转发端口状态: ${GREEN}已开放${NC}"
    else
        echo -e "转发端口状态: ${RED}未开放${NC}"
    fi
    
    # 检查隧道进程
    if pgrep -f "autossh.*${SERVER_PORT}" >/dev/null; then
        echo -e "隧道进程状态: ${GREEN}运行中${NC}"
    else
        echo -e "隧道进程状态: ${RED}未运行${NC}"
    fi
    
    echo -e "\n${GREEN}=== 连接步骤 ===${NC}"
    echo -e "1. 确保隧道显示"运行中""
    echo -e "2. 确保转发端口显示"已开放""
    echo -e "3. 使用以下命令连接："
    echo -e "   ${YELLOW}ssh -p ${SERVER_PORT} ${LOCAL_USER}@${SERVER_IP}${NC}"
    
    echo -e "\n${GREEN}=== 故障排除 ===${NC}"
    echo -e "1. 如果隧道未运行：选择菜单选项 1 启动隧道"
    echo -e "2. 如果端口未开放：检查服务器防火墙设置"
    echo -e "3. 如果配置有误：确保服务器 /etc/ssh/sshd_config 包含："
    echo -e "   ${YELLOW}GatewayPorts yes"
    echo -e "   AllowTcpForwarding yes${NC}"
    echo -e "4. 配置修改后需要重启SSH服务："
    echo -e "   ${YELLOW}sudo systemctl restart sshd${NC}"
    echo -e "5. 如果仍然无法连接：查看日志（选项 5）"
}

# 修改显示菜单函数
show_menu() {
    echo -e "\n${GREEN}=== 远程控制管理 ===${NC}"
    echo "1. 启动隧道"
    echo "2. 停止隧道"
    echo "3. 查看状态"
    echo "4. 修改配置"
    echo "5. 查看日志"
    echo "6. 安装系统服务"
    echo "7. 显示连接信息"
    echo "8. 修复音频系统"
    echo "0. 退出"
    echo -e "${YELLOW}请选择操作 [0-8]:${NC}"
}

# 添加音频系统修复函数
fix_audio_system() {
    echo -e "${YELLOW}正在修复音频系统...${NC}"
    
    # 安装必要的包
    echo -e "安装音频相关包..."
    sudo apt-get update
    sudo apt-get install -y pulseaudio pulseaudio-utils alsa-utils
    
    # 停止现有的 PulseAudio 进程
    echo -e "重启 PulseAudio 服务..."
    pulseaudio -k || true
    sleep 2
    
    # 删除旧的配置
    rm -rf ~/.config/pulse/* || true
    
    # 重新启动 PulseAudio
    pulseaudio --start
    sleep 2
    
    # 设置用户权限
    echo -e "设置用户权限..."
    sudo usermod -aG audio,pulse,pulse-access $USER
    
    # 创建/修改 PulseAudio 配置
    echo -e "配置 PulseAudio..."
    mkdir -p ~/.config/pulse
    cat > ~/.config/pulse/client.conf << EOF
autospawn = yes
daemon-binary = /usr/bin/pulseaudio
extra-arguments = --log-target=syslog
EOF

    # 重启 PulseAudio 守护进程
    pulseaudio -k
    sleep 2
    pulseaudio --start --log-target=syslog
    
    # 检查状态
    echo -e "\n${GREEN}=== 音频系统状态 ===${NC}"
    if pulseaudio --check; then
        echo -e "PulseAudio 状态: ${GREEN}运行中${NC}"
    else
        echo -e "PulseAudio 状态: ${RED}未运行${NC}"
    fi
    
    # 显示设备信息
    echo -e "\n${YELLOW}音频设备列表：${NC}"
    pactl list short sinks || true
    
    echo -e "\n${YELLOW}麦克风设备列表：${NC}"
    pactl list short sources || true
}

# 主程序
main() {
    check_dependencies
    load_config
    setup_ssh_key
    
    while true; do
        show_menu
        read -r choice
        
        case $choice in
            1)
                echo -e "${YELLOW}启动隧道...${NC}"
                create_tunnel
                echo -e "${GREEN}隧道已启动${NC}"
                ;;
            2)
                echo -e "${YELLOW}停止隧道...${NC}"
                pkill -f "autossh.*$SERVER_PORT"
                echo -e "${GREEN}隧道已停止${NC}"
                ;;
            3)
                if pgrep -f "autossh.*$SERVER_PORT" > /dev/null; then
                    echo -e "${GREEN}隧道状态: 运行中${NC}"
                else
                    echo -e "${RED}隧道状态: 已停止${NC}"
                fi
                ;;
            4)
                setup_config
                ;;
            5)
                echo -e "${YELLOW}最近的日志:${NC}"
                journalctl -u ssh-tunnel -n 50
                ;;
            6)
                install_service
                ;;
            7)
                show_connection_info
                ;;
            8)
                fix_audio_system
                ;;
            0)
                echo -e "${GREEN}再见！${NC}"
                exit 0
                ;;
            *)
                echo -e "${RED}无效的选择${NC}"
                ;;
        esac
    done
}

# 启动主程序
main